/*
 *  Copyright 2013 National Institute of Advanced Industrial Science and Technology
 *  
 *  Licensed under the Apache License, Version 2.0 (the "License");
 *  you may not use this file except in compliance with the License.
 *  You may obtain a copy of the License at
 *  
 *      http://www.apache.org/licenses/LICENSE-2.0
 *  
 *  Unless required by applicable law or agreed to in writing, software
 *  distributed under the License is distributed on an "AS IS" BASIS,
 *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 *  See the License for the specific language governing permissions and
 *  limitations under the License.
 */

package org.sss.common.io;

import java.io.IOException;
import java.util.ArrayDeque;
import java.util.ArrayList;
import java.util.Deque;
import java.util.List;
import java.util.Map;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.sss.common.DataReceiver;
import org.sss.common.KeyValues;
import org.sss.mapreduce.Configuration;
import org.sss.mapreduce.DB;
import org.sss.mapreduce.GroupID;
import org.sss.mapreduce.SssException;
import org.sss.mapreduce.StorageNode;
import org.sss.util.BytesComparator;

public class RestrictedWorkerFirstKeyValueReader extends AbstractKeyValueReader {
  private static BytesComparator CMP = new BytesComparator();
  @SuppressWarnings("unused")
  private static Logger logger = LoggerFactory.getLogger(RestrictedWorkerFirstKeyValueReader.class);

  public RestrictedWorkerFirstKeyValueReader(Configuration conf,
                                             StorageNode node,
                                             GroupID source,
                                             int subhash,
                                             DataReceiver<List<KeyValues>> next) throws SssException {
    super(conf, node, source, subhash, next);
  }

  @Override
  protected void doRead(IOThread iot)
    throws SssException, IOException, InterruptedException {
    try {
      DB db = iot.getDB(node);
      int count_max = iot.getIOGetBatchCount();
      int limitNBytes = iot.getIOGetLimitNBytes();
      WorkerQueue q = new WorkerQueue();
      collectWorkers_restricted(db, q);
      workerFirstReadRaw_restricted(db, limitNBytes, count_max, q);
    }
    catch (Exception e) {
      // TODO: error handling
      throw new IOException(e);
    }
  }
  private void collectWorkers_restricted(final DB db, WorkerQueue workerQueue)
    throws SssException, IOException {
    final KeyFormatter keyFormatter = KeyFormatter.get(getSource().getKeyFormat());
    assert keyFormatter == KeyFormatter.threadPriorityFormatter: "KeyFormatter is not threadPriorityFormatter.";
    final byte[] prefix = KeyFormatter.getPrefix(getSource(), subhash);
    final List<byte[]> allkeys = db.fwmkeys(prefix, -1);
    byte[] wid_ = new byte[]{}; // previous worker id
    Worker worker = null;
    for (int i=0; i<allkeys.size(); i++) {
      final byte[] key = allkeys.get(i);
      allkeys.set(i, null);                       // for gc
      byte[] wid = keyFormatter.getWorkerID(key); // worker id
      if (CMP.compare(wid, wid_) != 0) {
        if (worker != null) {
          workerQueue.push(worker.firstKey(), worker);
        }
        worker = new Worker(keyFormatter);
        worker.bytes = new ArrayList<byte[]>();
        wid_ = wid;
      }
      worker.bytes.add(key);
    }
    if (worker != null) {
      workerQueue.push(worker.firstKey(), worker);
    }
  }
  private void workerFirstReadRaw_restricted(DB db, int limitNBytes, int count_max,
      WorkerQueue workerQueue)
    throws InstantiationException, IllegalAccessException, Exception {
    final KeyFormatter keyFormatter = KeyFormatter.get(getSource().getKeyFormat());
    assert keyFormatter == KeyFormatter.threadPriorityFormatter: "KeyFormatter is not threadPriorityFormatter.";

    List<byte[]> ks = new ArrayList<byte[]>();
    List<byte[]> kvs = new ArrayList<byte[]>();
    Deque<Integer> akin = new ArrayDeque<Integer>();
    KeyValuesQueue kvq = new KeyValuesQueue(next, limitNBytes);
    try {
      while (!workerQueue.isEmpty()) {
        Map.Entry<byte[], List<Worker>> e = workerQueue.pop();
        byte[] key = e.getKey();
        List<Worker> workers = e.getValue();

        int sameUserKey = 0;
        for (Worker worker: workers) {
          for (; worker.index < worker.bytes.size(); worker.index++) {
            if (!keyFormatter.rawKeyEquals(worker.bytes.get(worker.index), key)) {
              workerQueue.push(worker.firstKey(), worker);
              break;
            }
            ks.add(worker.bytes.get(worker.index));
            worker.bytes.set(worker.index, null); // for gc
            sameUserKey++;
          }
        }
        akin.add(sameUserKey);
        if (ks.size() >= count_max || workerQueue.isEmpty()) {
          final int count_target = ks.size();
          boolean done = count_target == 0;
          int count_sofar = 0;
          while (!done) {
            checkCanceled();
            final int count_tmp = Math.min(count_target - count_sofar, count_max);
            byte[][] ks_ = new byte[count_tmp][];
            for (int i=0; i<count_tmp; i++) {
              ks_[i] = ks.get(count_sofar + i);
            }
            List<byte[]> kvs_ = db.getlist(ks_);
            countup(kvs_);
            kvs.addAll(kvs_);
            count_sofar += kvs_.size() / 2;
            done = count_target == count_sofar;
          }
          ks.clear();
        }

        while (akin.size() != 0 && akin.peek() <= kvs.size()/2) {
          kvq.pushKey(keyFormatter.getUserKey(kvs.get(0)));
          int n = akin.poll();
          for (int i = 0; i < n; ++i) {
            kvq.pushValue(kvs.get(i*2+1));
          }
          kvs = new ArrayList<byte[]>(kvs.subList(2 * n, kvs.size()));
        }
      }
    }
    finally {
      kvq.pushEnd();
    }
  }
}
